import numpy as np
from numpy.ma import cos
import matplotlib.pyplot as plt
from matplotlib import cm
import math
import time
import matplotlib

matplotlib.use('TkAgg')
matplotlib.rcParams['font.family'] = 'SimHei'
matplotlib.rcParams['axes.unicode_minus'] = False


class GeneticAlgorithmTester:
    def __init__(self):
        # 算法参数
        self.DNA_SIZE = 24  # 编码长度
        self.CROSS_RATE = 0.8  # 交叉率
        self.MUTA_RATE = 0.15  # 变异率
        self.Iterations = 50  # 迭代次数
        self.X_BOUND = [0, 10]  # X区间
        self.Y_BOUND = [0, 10]  # Y区间

        # 测试不同的种群规模
        self.POP_SIZES = [20, 50, 100, 200, 500]

        # 创建可视化图形
        self.fig = plt.figure(figsize=(18, 10))
        self.setup_visualization()

    def setup_visualization(self):
        """初始化可视化设置"""
        # 3D函数可视化
        self.ax_3d = self.fig.add_subplot(231, projection='3d')
        X = np.linspace(*self.X_BOUND, 100)
        Y = np.linspace(*self.Y_BOUND, 100)
        X, Y = np.meshgrid(X, Y)
        Z = self.fitness_function(X, Y)
        self.ax_3d.plot_surface(X, Y, Z, rstride=5, cstride=5, cmap=cm.coolwarm, alpha=0.6)
        self.ax_3d.set_xlabel('X')
        self.ax_3d.set_ylabel('Y')
        self.ax_3d.set_zlabel('Z')
        self.ax_3d.set_title('目标函数曲面')

        # 适应度曲线图
        self.ax_fitness = self.fig.add_subplot(232)
        self.ax_fitness.set_xlabel('迭代次数')
        self.ax_fitness.set_ylabel('适应度')
        self.ax_fitness.set_title('适应度变化曲线')
        self.ax_fitness.grid(True)

        # 运行时间比较图
        self.ax_time = self.fig.add_subplot(233)
        self.ax_time.set_xlabel('种群规模')
        self.ax_time.set_ylabel('运行时间(秒)')
        self.ax_time.set_title('运行时间比较')
        self.ax_time.grid(True)

        # 种群分布图
        self.ax_pop = self.fig.add_subplot(234)
        self.ax_pop.set_xlabel('X')
        self.ax_pop.set_ylabel('Y')
        self.ax_pop.set_title('最优解分布')
        self.ax_pop.set_xlim(self.X_BOUND)
        self.ax_pop.set_ylim(self.Y_BOUND)

        # 适应度比较图
        self.ax_compare = self.fig.add_subplot(235)
        self.ax_compare.set_xlabel('迭代次数')
        self.ax_compare.set_ylabel('适应度')
        self.ax_compare.set_title('不同种群规模适应度比较')
        self.ax_compare.grid(True)

        # 最佳解比较图
        self.ax_solution = self.fig.add_subplot(236)
        self.ax_solution.set_xlabel('种群规模')
        self.ax_solution.set_ylabel('适应度')
        self.ax_solution.set_title('最终适应度比较')
        self.ax_solution.grid(True)

        plt.tight_layout()

    def fitness_function(self, x, y):
        """适应度函数"""
        return (6.452 * (x + 0.125 * y) * (cos(x) - cos(2 * y)) ** 2) / (
                0.8 + (x - 4.2) ** 2 + 2 * (y - 7) ** 2) + 3.226 * y

    def decode_genome(self, population):
        """解码DNA序列为实际参数值"""
        x_pop = population[:, 1::2]  # 奇数列表示X
        y_pop = population[:, ::2]  # 偶数列表示y

        # 将二进制转换为实数
        x = x_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.X_BOUND[1] - self.X_BOUND[0]) + self.X_BOUND[0]
        y = y_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * \
            (self.Y_BOUND[1] - self.Y_BOUND[0]) + self.Y_BOUND[0]

        return x, y

    def calculate_fitness(self, population):
        """计算种群适应度"""
        x, y = self.decode_genome(population)
        values = self.fitness_function(x, y)
        return (values - np.min(values)) + 0.0001  # 防止适应度出现负数

    def selection(self, population, fitness):
        """选择操作 - 轮盘赌选择"""
        idx = np.random.choice(
            np.arange(population.shape[0]),
            size=population.shape[0],
            replace=True,
            p=fitness / fitness.sum()
        )
        return population[idx]

    def crossover_and_mutation(self, population):
        """交叉和变异操作"""
        new_population = []
        for individual in population:
            # 创建副本避免修改原个体
            child = individual.copy()

            # 交叉操作
            if np.random.rand() < self.CROSS_RATE:
                mate_idx = np.random.randint(population.shape[0])
                crossover_point1 = np.random.randint(0, self.DNA_SIZE * 2 - 1)
                crossover_point2 = np.random.randint(crossover_point1, self.DNA_SIZE * 2)
                child[crossover_point1:crossover_point2] = \
                    population[mate_idx][crossover_point1:crossover_point2]

            # 变异操作
            self.mutate(child)

            new_population.append(child)

        return np.array(new_population)

    def mutate(self, individual):
        """变异操作"""
        if np.random.rand() < self.MUTA_RATE:
            mutate_point = np.random.randint(0, len(individual))
            individual[mutate_point] = individual[mutate_point] ^ 1

    def run_optimization(self, pop_size):
        """运行遗传算法优化并记录时间"""
        start_time = time.time()

        # 初始化种群
        population = np.random.randint(2, size=(pop_size, self.DNA_SIZE * 2))
        best_fitness_history = []
        best_solutions = []

        for iteration in range(self.Iterations):
            # 遗传操作
            population = self.crossover_and_mutation(population)
            fitness = self.calculate_fitness(population)

            # 解码当前种群
            x, y = self.decode_genome(population)
            current_fitness = self.fitness_function(x, y)

            # 记录最佳适应度和解
            best_idx = np.argmax(current_fitness)
            best_fitness_history.append(current_fitness[best_idx])
            best_solutions.append((x[best_idx], y[best_idx]))

            # 选择新一代
            population = self.selection(population, fitness)

        runtime = time.time() - start_time

        return {
            'pop_size': pop_size,
            'best_fitness': best_fitness_history[-1],
            'runtime': runtime,
            'best_solution': best_solutions[-1],
            'fitness_history': best_fitness_history
        }

    def test_population_sizes(self):
        """测试不同种群规模"""
        all_results = []
        runtimes = []

        for pop_size in self.POP_SIZES:
            print(f"\n正在运行种群规模 {pop_size}...")
            result = self.run_optimization(pop_size)
            all_results.append(result)
            runtimes.append(result['runtime'])

            # 更新运行时间比较图
            self.ax_time.clear()
            self.ax_time.bar(range(len(self.POP_SIZES[:len(runtimes)])), runtimes,
                             tick_label=self.POP_SIZES[:len(runtimes)])
            self.ax_time.set_xlabel('种群规模')
            self.ax_time.set_ylabel('运行时间(秒)')
            self.ax_time.set_title('运行时间比较')
            self.ax_time.grid(True)

            # 更新适应度比较图
            self.update_comparison_charts(all_results)

            # 显示当前结果
            print(f"种群规模 {pop_size}:")
            print(f"最终适应度: {result['best_fitness']:.4f}")
            print(f"运行时间: {result['runtime']:.2f}秒")
            print(f"最佳解: {result['best_solution']}")

            plt.pause(0.1)

        # 最终比较
        self.final_comparison(all_results)
        plt.show()

    def update_comparison_charts(self, all_results):
        """更新比较图表"""
        # 适应度比较
        self.ax_compare.clear()
        for result in all_results:
            self.ax_compare.plot(result['fitness_history'], label=f'种群 {result["pop_size"]}')
        self.ax_compare.set_xlabel('迭代次数')
        self.ax_compare.set_ylabel('适应度')
        self.ax_compare.set_title('不同种群规模适应度比较')
        self.ax_compare.legend()
        self.ax_compare.grid(True)

        plt.draw()

    def final_comparison(self, all_results):
        """绘制最终比较图"""
        # 最佳适应度比较
        self.ax_solution.clear()
        pop_sizes = [r['pop_size'] for r in all_results]
        final_fitness = [r['best_fitness'] for r in all_results]
        self.ax_solution.bar(range(len(pop_sizes)), final_fitness,
                             tick_label=pop_sizes)
        self.ax_solution.set_xlabel('种群规模')
        self.ax_solution.set_ylabel('最终适应度')
        self.ax_solution.set_title('不同种群规模的最终适应度')
        self.ax_solution.grid(True)

        # 解分布比较
        self.ax_pop.clear()
        for result in all_results:
            x, y = result['best_solution']
            self.ax_pop.scatter(x, y, s=result['pop_size'] / 5,
                                label=f'种群 {result["pop_size"]}', alpha=0.6)
        self.ax_pop.set_xlim(self.X_BOUND)
        self.ax_pop.set_ylim(self.Y_BOUND)
        self.ax_pop.set_title('不同种群规模的最优解分布')
        self.ax_pop.legend()

        plt.tight_layout()


if __name__ == "__main__":
    tester = GeneticAlgorithmTester()
    tester.test_population_sizes()